# data_sources/baostock_source.py
import baostock as bs
import pandas as pd
import numpy as np
from .base_data_source import BaseDataSource
import logging
from datetime import datetime, timedelta

class BaostockDataSource(BaseDataSource):
    def __init__(self):
        super().__init__("baostock")
        self.lg = None
        self._login()

    def _login(self):
        """登录Baostock"""
        try:
            self.lg = bs.login()
            if self.lg.error_code == '0':
                logging.info("Baostock登录成功")
            else:
                logging.error(f"Baostock登录失败: {self.lg.error_msg}")
        except Exception as e:
            logging.error(f"Baostock登录异常: {e}")

    def __del__(self):
        """退出时自动登出"""
        if self.lg:
            bs.logout()

    def get_stock_list(self, market: str = "A") -> pd.DataFrame:
        """
        获取股票列表
        market: A-全部, sz-深圳, sh-上海
        """
        try:
            if not self.lg:
                self._login()

            # 获取证券基本资料
            rs = bs.query_stock_basic(market)
            if rs.error_code != '0':
                logging.error(f"获取股票列表失败: {rs.error_msg}")
                return pd.DataFrame()

            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())

            df = pd.DataFrame(data_list, columns=rs.fields)

            # 标准化列名
            column_mapping = {
                'code': 'symbol',
                'code_name': 'name',
                'ipoDate': 'listing_date',
                'outDate': 'delisting_date',
                'type': 'type'
            }
            df = df.rename(columns=column_mapping)

            # 添加市场信息
            df['market'] = df['symbol'].apply(self._get_market_from_symbol)
            df['country'] = 'CN'
            df['currency'] = 'CNY'

            # 过滤掉已退市股票
            df = df[df['status'] == '1']  # 1-上市，0-退市

            logging.info(f"从Baostock获取到 {len(df)} 只股票信息")
            return df

        except Exception as e:
            logging.error(f"获取Baostock股票列表异常: {e}")
            return pd.DataFrame()

    def _get_market_from_symbol(self, symbol: str) -> str:
        """根据股票代码判断市场"""
        if symbol.startswith('sh.') or symbol.startswith('sz.'):
            code = symbol[3:]
        else:
            code = symbol

        if code.startswith('6'):
            return 'SH'
        elif code.startswith('0') or code.startswith('3'):
            return 'SZ'
        elif code.startswith('4') or code.startswith('8'):
            return 'BJ'  # 北京交易所
        else:
            return 'UNKNOWN'

    def get_daily_data(self, symbol: str, start_date: str, end_date: str) -> pd.DataFrame:
        """
        获取日线数据
        symbol格式: sh.600000 或 sz.000001
        """
        try:
            if not self.lg:
                self._login()

            # 处理symbol格式
            if not symbol.startswith(('sh.', 'sz.')):
                symbol = self._format_symbol(symbol)

            # 调整日期格式
            start_date = start_date.replace('-', '')
            end_date = end_date.replace('-', '')

            # 查询复权因子
            rs_factor = bs.query_adjust_factor(
                code=symbol,
                start_date=start_date,
                end_date=end_date
            )

            factor_data = []
            if rs_factor.error_code == '0':
                while (rs_factor.error_code == '0') & rs_factor.next():
                    factor_data.append(rs_factor.get_row_data())

            factor_df = pd.DataFrame(factor_data, columns=rs_factor.fields)

            # 查询历史K线数据
            rs = bs.query_history_k_data_plus(
                symbol,
                "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST",
                start_date=start_date,
                end_date=end_date,
                frequency="d",
                adjustflag="3"  # 3-后复权
            )

            if rs.error_code != '0':
                logging.error(f"获取日线数据失败 {symbol}: {rs.error_msg}")
                return pd.DataFrame()

            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())

            if not data_list:
                return pd.DataFrame()

            df = pd.DataFrame(data_list, columns=rs.fields)

            # 数据类型转换
            numeric_columns = ['open', 'high', 'low', 'close', 'preclose', 'volume',
                               'amount', 'turn', 'pctChg', 'peTTM', 'pbMRQ', 'psTTM', 'pcfNcfTTM']
            for col in numeric_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

            # 标准化列名
            df = self.normalize_daily_data(df)

            # 添加复权因子信息
            if not factor_df.empty:
                factor_df = factor_df.rename(columns={'code': 'symbol', 'exdiv_date': 'ex_dividend_date'})
                # 这里可以根据需要合并复权因子数据

            logging.info(f"从Baostock获取 {symbol} 日线数据 {len(df)} 条")
            return df

        except Exception as e:
            logging.error(f"获取Baostock日线数据异常 {symbol}: {e}")
            return pd.DataFrame()

    def get_financial_data(self, symbol: str, year: int = None, quarter: int = None) -> pd.DataFrame:
        """
        获取财务数据
        """
        try:
            if not self.lg:
                self._login()

            # 处理symbol格式
            if not symbol.startswith(('sh.', 'sz.')):
                symbol = self._format_symbol(symbol)

            # 查询成长能力数据
            growth_list = []
            rs_growth = bs.query_growth_data(
                code=symbol,
                year=year,
                quarter=quarter
            )
            if rs_growth.error_code == '0':
                while (rs_growth.error_code == '0') & rs_growth.next():
                    growth_list.append(rs_growth.get_row_data())

            # 查询盈利能力数据
            profit_list = []
            rs_profit = bs.query_profit_data(
                code=symbol,
                year=year,
                quarter=quarter
            )
            if rs_profit.error_code == '0':
                while (rs_profit.error_code == '0') & rs_profit.next():
                    profit_list.append(rs_profit.get_row_data())

            # 查询营运能力数据
            operation_list = []
            rs_operation = bs.query_operation_data(
                code=symbol,
                year=year,
                quarter=quarter
            )
            if rs_operation.error_code == '0':
                while (rs_operation.error_code == '0') & rs_operation.next():
                    operation_list.append(rs_operation.get_row_data())

            # 合并财务数据（这里简化处理，实际可以根据需要详细处理）
            financial_data = {}

            if growth_list:
                growth_df = pd.DataFrame(growth_list, columns=rs_growth.fields)
                financial_data.update(self._parse_growth_data(growth_df))

            if profit_list:
                profit_df = pd.DataFrame(profit_list, columns=rs_profit.fields)
                financial_data.update(self._parse_profit_data(profit_df))

            financial_df = pd.DataFrame([financial_data])
            financial_df['symbol'] = symbol
            financial_df['data_source'] = self.source_name

            return financial_df

        except Exception as e:
            logging.error(f"获取Baostock财务数据异常 {symbol}: {e}")
            return pd.DataFrame()

    def _parse_growth_data(self, df: pd.DataFrame) -> dict:
        """解析成长能力数据"""
        if df.empty:
            return {}

        data = {}
        latest = df.iloc[0]  # 取最新数据

        mappings = {
            'YOYEquity': 'equity_growth_rate',
            'YOYAsset': 'asset_growth_rate',
            'YOYNI': 'net_profit_growth_rate',
            'YOYEPSBasic': 'eps_growth_rate'
        }

        for bs_key, our_key in mappings.items():
            if bs_key in latest:
                data[our_key] = latest[bs_key]

        return data

    def _parse_profit_data(self, df: pd.DataFrame) -> dict:
        """解析盈利能力数据"""
        if df.empty:
            return {}

        data = {}
        latest = df.iloc[0]

        mappings = {
            'roeAvg': 'roe',
            'npMargin': 'net_margin',
            'gpMargin': 'gross_margin',
            'netProfit': 'net_profit',
            'epsTTM': 'eps',
            'MBRevenue': 'revenue'
        }

        for bs_key, our_key in mappings.items():
            if bs_key in latest:
                data[our_key] = latest[bs_key]

        return data

    def _format_symbol(self, symbol: str) -> str:
        """格式化股票代码"""
        if symbol.startswith('6'):
            return f"sh.{symbol}"
        elif symbol.startswith('0') or symbol.startswith('3'):
            return f"sz.{symbol}"
        elif symbol.startswith('4') or symbol.startswith('8'):
            return f"bj.{symbol}"  # 北京交易所
        else:
            return symbol

    def get_dividend_data(self, symbol: str, year: str = None) -> pd.DataFrame:
        """
        获取分红数据
        """
        try:
            if not self.lg:
                self._login()

            symbol = self._format_symbol(symbol)

            rs = bs.query_dividend_data(
                code=symbol,
                year=year,
                yearType="report"
            )

            if rs.error_code != '0':
                logging.error(f"获取分红数据失败 {symbol}: {rs.error_msg}")
                return pd.DataFrame()

            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())

            if not data_list:
                return pd.DataFrame()

            df = pd.DataFrame(data_list, columns=rs.fields)

            # 标准化列名
            column_mapping = {
                'code': 'symbol',
                'dividendDate': 'dividend_date',
                'dividend': 'dividend_per_share',
                'ratio': 'dividend_ratio'
            }
            df = df.rename(columns=column_mapping)

            # 数据类型转换
            numeric_cols = ['dividend_per_share', 'dividend_ratio']
            for col in numeric_cols:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

            return df

        except Exception as e:
            logging.error(f"获取Baostock分红数据异常 {symbol}: {e}")
            return pd.DataFrame()

    def normalize_daily_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """重写标准化方法，针对Baostock的特殊字段"""
        if df.empty:
            return df

        # Baostock特定列名映射
        baostock_mapping = {
            'date': 'trade_date',
            'code': 'symbol',
            'open': 'open_price',
            'high': 'high_price',
            'low': 'low_price',
            'close': 'close_price',
            'volume': 'volume',
            'amount': 'turnover',
            'pctChg': 'change_rate',
            'turn': 'turnover_rate',
            'peTTM': 'pe_ratio',
            'pbMRQ': 'pb_ratio'
        }

        df = df.rename(columns=baostock_mapping)

        # 确保必要的列存在
        required_columns = ['symbol', 'trade_date', 'open_price', 'high_price',
                            'low_price', 'close_price', 'volume']

        for col in required_columns:
            if col not in df.columns:
                df[col] = np.nan

        return df